# @Author: Zhixuan.Wang
# @IDE: PyCharm
# @Project: multimodal
# @File: MyDataSet.py
# @Time: 2025/11/6 11:09
# @Description: 构建自己的数据类

from torch.utils.data import Dataset
import os
import xml.etree.ElementTree as ET
from PIL import Image
import torch
import pandas as pd
import numpy as np


class MyImgDataset(Dataset):

    # The Initilization of Function
    def __init__(self, img_path, xml_path, transform=None):
        """
            img_path: 图片的路径 D:\LungCancer\VOCdevkit\VOC2012\EffectiveImage
            transform: 数据增强
            xml_path: 标签数据路径 D:\LungCancer\VOCdevkit\VOC2012\Annotations
        """
        self.img_path = img_path
        self.transform = transform
        self.xml_path = xml_path

        self.samples = []

        # 获取文件列表，过滤非图片文件
        self.imgs = [f for f in os.listdir(self.img_path)
                     if f.lower().endswith(('.jpg', '.png', '.jpeg', '.bmp'))]
        self.xmls = [f for f in os.listdir(self.xml_path) if f.endswith('.xml')]

        # 按文件名匹配图片和XML
        for img_file in self.imgs:
            img_name = os.path.splitext(img_file)[0]
            xml_file = f"{img_name}.xml"

            single_pic = os.path.join(self.img_path, img_file)
            single_xml = os.path.join(self.xml_path, xml_file)

            try:
                tree = ET.parse(single_xml)
                root = tree.getroot()

                for obj in root.findall('object'):
                    class_name = obj.find('name').text
                    if class_name == '良性结节':
                        class_label = 1
                    elif class_name == '恶性结节':
                        class_label = 2
                    else:
                        class_label = 0
                        print(f"警告：未知类别 {class_name}，标记为0")

                    bndbox = obj.find('bndbox')
                    xmin = int(bndbox.find('xmin').text)
                    ymin = int(bndbox.find('ymin').text)
                    xmax = int(bndbox.find('xmax').text)
                    ymax = int(bndbox.find('ymax').text)

                    position_label = {'xmin': xmin, 'ymin': ymin, 'xmax': xmax, 'ymax': ymax}

                    self.samples.append({
                        'single_pic': single_pic,
                        'class_label': class_label,
                        'position_label': position_label
                    })

            except Exception as e:
                print(f"解析文件 {single_xml} 时出错: {e}")

    # Get the Length of Data
    def __len__(self):
        return len(self.samples)

    # Get the Image due to certain ID
    def __getitem__(self, idx):
        sample = self.samples[idx]
        # 加载单个图片
        img_path = sample['single_pic']
        image = Image.open(img_path)

        #  转换为numpy再转tensor
        image = np.array(image)
        # 增加通道数，符合pytorch的输入格式
        image = torch.from_numpy(image).float().unsqueeze(0)  # [1, H, W]

        # 处理标签
        class_label = torch.tensor(sample['class_label'], dtype=torch.long)

        # 处理边界框
        pos = sample['position_label']
        bbox = torch.tensor([pos['xmin'], pos['ymin'], pos['xmax'], pos['ymax']], dtype=torch.float)

        return image, class_label, bbox
